# train4.py
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from model import get_model
# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # 将图像大小调整为固定的 256x256
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 归一化
])
# 加载自定义数据集
# 使用苹果与西红柿数据集，包含"苹果"和"番茄"两个子文件夹
train_dataset = datasets.ImageFolder(root='./苹果与西红柿数据集', transform=transform)
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
# 获取模型
model = get_model(num_classes=2)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 指定设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)


# 训练模型
def train(model, device, train_loader, optimizer, criterion, epochs):
    model.train()
    best_loss = float('inf')  # 初始化最低损失值为无穷大
    model_path = '模型文件.pth'  # 模型保存路径
    for epoch in range(epochs):
        epoch_loss = 0.0  # 初始化当前 epoch 的总损失
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()  # 累加当前 epoch 的损失
            if batch_idx % 100 == 0:
                print(
                    f'模型训练第 {epoch + 1} 批：[{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
        avg_epoch_loss = epoch_loss / len(train_loader)  # 计算当前 epoch 的平均损失
        print(f"Epoch {epoch + 1} 结束，平均损失: {avg_epoch_loss:.6f}")
        # 检查当前 epoch 的平均损失是否低于最低损失
        if avg_epoch_loss < best_loss:
            best_loss = avg_epoch_loss
            # 保存模型
            torch.save(model, model_path)
            print(f"模型已保存到 {model_path}，当前最低损失: {best_loss:.3f}")
    print("训练结束，模型文件保存成功")


# 训练和测试
train(model, device, train_loader, optimizer, criterion, epochs=50)
